#include "segment.h"
#include "helper.h"
#include "rapidjson/stringbuffer.h"
#include "rapidjson/writer.h"

#include <malloc.h>

Segment::Segment(const std::string &serviceId, const std::string &serviceInstanceId, int version, const std::string &header)
{

    _serviceId = serviceId;
    _serviceInstanceId = serviceInstanceId;
    _header = header;
    _version = version;

    std::string tmp_traceId = uuid() + "." + std::to_string(getpid()) + "." + std::to_string(random(999999));
    std::string traceId = str_replace(tmp_traceId, "-", "");

    Segment::_traceId = traceId;
    Segment::_traceSegmentId = traceId;
    bag = new CrossProcessBag(_serviceId, _serviceInstanceId, traceId, _version, _header);

    if (!bag->getTraceId().empty())
    {
        Segment::_traceId = bag->getTraceId();
    }
}

Segment::~Segment()
{
    delete bag;
    spans.clear();
    // std::vector<Span *>().swap(spans);
    spans.shrink_to_fit();
    malloc_trim(0);
}

std::string Segment::marshal()
{
    if (!spans.empty())
    {
        auto span = spans.front();
        span->setEndTime();
        if (_status_code >= 400)
        {
            span->setIsError(true);
        }
        span->pushTag(new Tag("status_code", std::to_string(_status_code)));
    }

    rapidjson::StringBuffer strBuf;
    rapidjson::Writer<rapidjson::StringBuffer> writer(strBuf);

    writer.StartObject();

    writer.Key("traceId");
    writer.String(_traceId.c_str());

    writer.Key("traceSegmentId");
    writer.String(_traceSegmentId.c_str());

    writer.Key("service");
    writer.String(_serviceId.c_str());

    writer.Key("serviceInstance");
    writer.String(_serviceInstanceId.c_str());

    writer.Key("isSizeLimited");
    writer.Bool(false);

    writer.Key("spans");
    writer.StartArray();
    for (auto span : spans)
    {
        writer.StartObject();

        writer.Key("spanId");
        writer.Int(span->getSpanId());

        writer.Key("parentSpanId");
        writer.Int(span->getParentSpanId());

        writer.Key("startTime");
        writer.Int64(span->getStartTime());

        writer.Key("endTime");
        writer.Int64(span->getEndTime());

        writer.Key("operationName");
        writer.String(span->getOperationName().c_str());

        writer.Key("peer");
        writer.String(span->getPeer().c_str());

        int span_type = (int)span->getSpanType();
        writer.Key("spanType");
        writer.Int(span_type);

        int span_layer = (int)span->getSpanLayer();
        writer.Key("spanLayer");
        writer.Int(span_layer);

        writer.Key("componentId");
        writer.Int(span->getComponentId());

        writer.Key("isError");
        writer.Bool(span->getIsError());

        writer.Key("refs");
        writer.StartArray();
        for (auto ref : span->getRefs())
        {
            writer.StartObject();

            writer.Key("refType");
            writer.Int(ref->getRefType());

            writer.Key("traceId");
            writer.String(ref->getTraceId().c_str());

            writer.Key("parentTraceSegmentId");
            writer.String(ref->getParentTraceSegmentId().c_str());

            writer.Key("parentSpanId");
            writer.Int(ref->getParentSpanId());

            writer.Key("parentService");
            writer.String(ref->getParentService().c_str());

            writer.Key("parentServiceInstance");
            writer.String(ref->getParentServiceInstance().c_str());

            writer.Key("parentEndpoint");
            writer.String(ref->getParentEndpoint().c_str());

            writer.Key("networkAddressUsedAtPeer");
            writer.String(ref->getNetworkAddressUsedAtPeer().c_str());

            writer.EndObject();
        }
        writer.EndArray();

        writer.Key("tags");
        writer.StartArray();
        for (auto tag : span->getTags())
        {
            writer.StartObject();

            writer.Key("key");
            writer.String(tag->getKey().c_str());

            writer.Key("value");
            writer.String(tag->getValue().c_str());

            writer.EndObject();
        }
        writer.EndArray();

        writer.EndObject();
    }
    writer.EndArray();
    writer.EndObject();

    std::string res = strBuf.GetString();
    
    return res;
}

void Segment::setStatusCode(int code)
{
    _status_code = code;
}

Span *Segment::createSpan(SkySpanType type, SkySpanLayer layer, int componentId)
{
    Span *span = new Span();
    span->setSpanType(type);
    span->setSpanLayer(layer);
    span->setComponentId(componentId);

    int id = 0;
    int parentId = -1;
    if (!spans.empty())
    {
        id = spans.back()->getSpanId() + 1;
        parentId = 0;
    }

    span->setSpanId(id);
    span->setParentSpanId(parentId);

    spans.push_back(span);
    return span;
}

Span *Segment::getLastSpan()
{
    return spans.back();
}

std::string Segment::createHeader(Span *span)
{
    return bag->encode(span->getSpanId(), span->getPeer());
}

void Segment::createRefs()
{
    if (!spans.empty())
    {
        Span *sp = spans.front();
        bag->setOperationName(sp->getOperationName());

        if (!_header.empty() && !bag->getParentTraceSegmentId().empty())
        {
            auto ref = new SkySegmentReference();
            ref->setRefType(0);
            ref->setTraceId(bag->getTraceId());
            ref->setParentTraceSegmentId(bag->getParentTraceSegmentId());
            ref->setParentSpanId(bag->getParentSpanId());
            ref->setParentService(bag->getParentService());
            ref->setParentServiceInstance(bag->getParentServiceInstance());
            ref->setParentEndpoint(bag->getParentEndpoint());
            ref->setNetworkAddressUsedAtPeer(bag->getNetworkAddressUsedAtPeer());
            sp->pushRefs(ref);
        }
    }
}

std::string Segment::getTraceId()
{
    return _traceId;
}